#include "Scheduler.hpp"
#include "Fiber.hpp"
#include "Thread.hpp"
#include "../base/BaseLog.hpp"

#include <cassert>
#include <cstddef>
#include <memory>
#include <mutex>
#include <string>
#include <unistd.h>
#include <vector>

using namespace DBServer::coroutine;

/// @brief 当前线程的调度器
static thread_local Scheduler * sche = nullptr;
/// @brief 调度器的调度协程
static thread_local Fiber *sche_fiber = nullptr;

Scheduler::Scheduler(std::size_t threads, bool caller, const std::string& name_) noexcept
:   name(name_) {
    assert(threads > 0);

    // 如果caller为true，则表示调用线程也作为线程池中的一部分
    if (caller) 
    {
        Fiber::GetThis();
        --threads;

        assert(GetThis() == nullptr);
        call_fiber = std::make_shared<Fiber>([this](){ this->run(); });

        sche = this;
        sche_fiber = call_fiber.get();
        main_thread = Thread::ThreadId();
    }
    else
    {
        // caller为false， 则表示调用线程不参与调度
        main_thread = -1;
        call_fiber = nullptr;
    }

    t_count = threads;
}

Scheduler::~Scheduler() noexcept {
    if (GetThis() == this)
    {
        sche = nullptr;
    }
}

Scheduler* Scheduler::GetThis() {
    return sche;
}

Fiber* Scheduler::GetMainFiber() {
    return sche_fiber;
}

void Scheduler::start() {
    std::unique_lock<std::mutex> lock(mtx);

    if (!pause) {
        warn("scheduler(%s) is already working.\n", name.c_str());
        return;
    }
    pause = false;

    threads.resize(t_count);
    t_ids.resize(t_count);

    try
    {
        for (int i = 0; i < t_count; ++ i)
        {
            threads[i] = Thread::CreatThread([this, i]() {
                // 初始化当前线程
                Thread::MainThread();
                // 设置线程id
                this->t_ids[i] = Thread::ThreadId();
                // 执行调度函数
                this->run();
            });
        }
    }
    catch (const std::exception& e)
    {
        error(e.what());
        throw;
    }
    debug("scheduler(%s) running.\n", name.c_str());
}

void Scheduler::stop() {
    if (stopping()) 
    {
        return;
    }

    {    
        std::unique_lock<std::mutex> lock(mtx);
        pause = true;

        if (main_thread != -1)
        {
            assert(GetThis() == this);
        }
        else
        {
            assert(GetThis() == nullptr);
        }

        // 向其他的线程发出信号唤醒正在等待的线程
        for (size_t i = 0; i < t_count; ++ i)
        {
            tickle();
        }

        // 如果调度器线程也参与了调度，
        // 就需要在发出一次信号
        if (call_fiber)
        {
            tickle();
        }
    }

    // 切换到调度器线程的调度协程
    if (call_fiber)
    {
        call_fiber->resume();
    }

    // 等待其他线程的结束
    for (auto &t : threads)
    {
        t.join();
    }
    info("scheduler(%s) stop!\n", name.c_str());
}

void Scheduler::switch_to(int thread) {
    assert(Scheduler::GetThis() != nullptr);
    if (Scheduler::GetThis() == this) 
    {
        if (thread == -1 || thread == Thread::ThreadId())
        {
            return;
        }
    }

    schedule(Fiber::GetThis(), thread);
    Fiber::Yield();
}

void Scheduler::tickle() {
    debug("tickle\n");
}

bool Scheduler::stopping() {
    std::unique_lock<std::mutex> lock(mtx);
    return pause &&
           tasks.empty() &&
           active_count == 0;
}

void Scheduler::run() {
    // 设置当前线程的调度器
    set_this();
    sche_fiber = Fiber::GetThis().get();

    Fiber::ptr idle_fiber = std::make_shared<Fiber>([this]() {
        this->idle();
    }, true);

    Fiber::ptr cb_fiber = nullptr;
    FiberAndThread ft;

    while (true)
    {
        ft.reset();
        bool tick = false;
        bool is_active = false;

        {
            std::unique_lock<std::mutex> lock(mtx);
            // 找到一个可执行的任务
            auto iter = tasks.begin();
            while (iter != tasks.end())
            {
                if (iter->t_id != -1 && iter->t_id != Thread::ThreadId())
                {
                    tick |= true;
                    ++iter;
                    continue;
                }

                assert(iter->cb || iter->fiber);
                // 当前协程正在执行中，调度器中只有一个任务队列
                // 所以需要进行判断
                if (iter->fiber && iter->fiber->get_state() == FiberState::EXEC)
                {
                    ++iter;
                    continue;
                }

                ft = std::move(*iter);
                tasks.erase(iter);
                is_active = true;
                ++active_count;
                ++iter;
                break;
            }
            // 如果任务队列中还有任务，需要唤醒其他的线程
            tick |= iter != tasks.end();
        }

        if (tick)
        {
            tickle();
        }

        if (ft.fiber && ft.fiber->get_state() != FiberState::TERM)
        {
            ft.fiber->resume();
            --active_count;

            if (ft.fiber->get_state() == FiberState::READY) {
                schedule(ft.fiber);
            } else if (ft.fiber->get_state() != FiberState::TERM) {
                ft.fiber->set_state(FiberState::HOLD);
            }
            ft.reset();
        }
        else if (ft.cb) {
            if (cb_fiber)  {
                cb_fiber->reset(ft.cb, true);
            } else {
                cb_fiber = std::make_shared<Fiber>(ft.cb, true);
            }
            cb_fiber->resume();
            --active_count;
            ft.reset();

            if (cb_fiber->get_state() != FiberState::TERM) {
                cb_fiber->set_state(FiberState::HOLD);
                cb_fiber.reset();
            } else { // cb_fiber->get_state() == FiberState::TERM
                cb_fiber->reset(nullptr);
            }
        }
        else
        {
            // 拿到的任务不可执行
            if (is_active) {
                --active_count;
                continue;
            }
            // 调度器停止工作
            if (idle_fiber->get_state() == FiberState::TERM) {
                info("scheduler thread(%d) term.\n", Thread::ThreadId());
                cb_fiber->set_state(FiberState::TERM);
                break;
            }
            ++idle_count;
            idle_fiber->resume();
            if (idle_fiber->get_state() != FiberState::TERM) {
                idle_fiber->set_state(FiberState::READY);
            }
            --idle_count;
        }
    }
}

void Scheduler::idle() {
    while (!stopping())
    {
        Fiber::Yield();
    }
}

void Scheduler::set_this() {
    sche = this;
}